{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Flash_lab",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# **嗨！欢迎来到小闪电的python实验室**\n",
        "⚡️⚡️⚡️⚡️⚡️⚡️⚡️⚡️"
      ],
      "metadata": {
        "id": "hH5cjFSySs-L"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L7aYhhBuCexz"
      },
      "source": [
        "# **跳甲运动轨迹跟踪及速度监测**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VOUH6SB09I2M"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "#挂载谷歌云盘"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y3p5a7V4EIAS"
      },
      "source": [
        "!pip install pims"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zSlG-p4QENpj"
      },
      "source": [
        "!pip install trackpy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8NwVppZ3D5EW"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from pandas import DataFrame, Series\n",
        "import pims\n",
        "import trackpy as tp\n",
        "import cv2\n",
        "import glob\n",
        "from PIL import Image\n",
        "import seaborn as sns"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RIA38ZmOD5Ea"
      },
      "source": [
        "# Master script\n",
        "\n",
        "type_ = input('Is is this a movie or folder of images? Please answer movie or images :' ).lower()\n",
        "\n",
        "if type_ == 'movie':\n",
        "            \n",
        "    # Import movie and get frames\n",
        "    vidcap = cv2.VideoCapture('/content/drive/MyDrive/track/ck/CK.mp4') # Define specific folder\n",
        "    success,image1 = vidcap.read()\n",
        "    image_list = []\n",
        "    count = 0\n",
        "\n",
        "    \n",
        "    # Select color channel\n",
        "    color = input('What color channel are you analyzing? (red, green, blue, all)')\n",
        "    color = color.lower()\n",
        "\n",
        "    if color == 'red':\n",
        "        color = 'r'\n",
        "    elif color == 'green':\n",
        "        color = 'g'\n",
        "    elif color == 'blue':\n",
        "        color = 'b'\n",
        "    elif color == 'all':\n",
        "        color = 'all'\n",
        "    else:\n",
        "        raise Exception(f\"color should be red(r), green(g), blue(b), or all. The value of color was {chann_color}\")\n",
        "    print(color)\n",
        "\n",
        "\n",
        "    # set threshold\n",
        "    thresh = int(input(\"Select starting threshold for creating white and black image (cells will be white, range 0-255)\"))\n",
        "\n",
        "    #If red channel, select red images\n",
        "    if color == 'r':\n",
        "        image = image1[:,:,2]\n",
        "        (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "        \n",
        "    #If red channel, select red images\n",
        "    elif color == 'g':\n",
        "        image = image1[:,:,1]\n",
        "        (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "    #If blue channel, select red images\n",
        "    elif color == 'b':\n",
        "        image = image1[:,:,0]\n",
        "        (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "    #If all channels, select red images\n",
        "    elif color == 'all':\n",
        "        image = image1[:,:]\n",
        "        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
        "        (_, bw) = cv2.threshold(gray, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "            \n",
        "    else:\n",
        "        raise Exception(f\"color should be red(r), green(g), blue(b), or all. The value of color was {color}\")\n",
        "\n",
        "    plt.figure()\n",
        "    plt.imshow(bw, plt.get_cmap('gray'))\n",
        "    plt.show()\n",
        "\n",
        "    ok = input('Is this threshold OK? (yes or no)')\n",
        "    ok = ok.lower()\n",
        "\n",
        "    while ok != 'yes':\n",
        "        try:\n",
        "            # set threshold\n",
        "            thresh = int(input(\"Select starting threshold for creating white and black image (cells will be white, range 0-255)\"))\n",
        "\n",
        "            #If red channel, select red images\n",
        "            if color == 'r':\n",
        "                image = image1[:,:,2]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "            #If red channel, select red images\n",
        "            elif color == 'g':\n",
        "                image = image1[:,:,1]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "            #If blue channel, select red images\n",
        "            elif color == 'b':\n",
        "                image = image1[:,:,0]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "            #If all channels, select red images\n",
        "            elif color == 'all':\n",
        "                image = image1[:,:]\n",
        "                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
        "                (_, bw) = cv2.threshold(gray, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "            plt.figure()\n",
        "            plt.imshow(bw, plt.get_cmap('gray'))\n",
        "            plt.show()\n",
        "\n",
        "            ok = input('Is this threshold OK? (yes or no)')\n",
        "            ok = ok.lower()\n",
        "\n",
        "        except:\n",
        "            raise Exception(\"Something went wrong while re-trying threshold\")\n",
        "            \n",
        "    #Append because vidcap will call the next image when called again\n",
        "    image_list.append(bw)\n",
        "\n",
        "    print('good job')\n",
        "\n",
        "    while success:\n",
        "        success,image = vidcap.read()\n",
        "        count += 1\n",
        "        try:\n",
        "         #If red channel, select red images\n",
        "            if color == 'r':\n",
        "                image = image[:,:,2]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "        \n",
        "            #If red channel, select red images\n",
        "            elif color == 'g':\n",
        "                image = image[:,:,1]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "            \n",
        "            #If blue channel, select red images\n",
        "            elif color == 'b':\n",
        "                image = image[:,:,0]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "            \n",
        "            #If all channels, select red images\n",
        "            elif color == 'all':\n",
        "                image = image[:,:]\n",
        "                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
        "                (_, bw) = cv2.threshold(gray, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "            \n",
        "            else:\n",
        "                raise Exception(f\"color should be red(r), green(g), blue(b), or all. The value of color was {color}\")\n",
        "            \n",
        "        except:\n",
        "            print(f\"No image was grabbed for image {count}\")\n",
        "            break\n",
        "\n",
        "    image_series = input('Would you like to see the first 9 frames? :').lower()\n",
        "    \n",
        "    if image_series == 'yes':\n",
        "        # Plot subplots of figures\n",
        "        plt.figure(figsize=(10,10))\n",
        "        for i in range(9):\n",
        "            # define subplot\n",
        "            plt.subplot(330 + 1 + i)\n",
        "            # convert to unsigned integers for viewing\n",
        "            image = image_list[i].astype('uint8')\n",
        "            # plot raw pixel data\n",
        "            plt.imshow(image, plt.get_cmap('gray'))\n",
        "\n",
        "if type_ == 'images':\n",
        "    \n",
        "    #Get images\n",
        "    image_list=[]\n",
        "    images = []\n",
        "    files = glob.glob(\"orig_movie/*.png\") # Define specific file path\n",
        "    names = np.asanyarray([file for file in files])\n",
        "    file_num = np.reshape(np.asanyarray([int(''.join(filter(str.isdigit, name))) for name in names]),(-1,1))\n",
        "    names = names.reshape((-1,1))\n",
        "    name_files = np.hstack((names,file_num))\n",
        "    name_files = name_files[np.argsort(name_files[:,1])]\n",
        "    for file in name_files[:,0]:\n",
        "        image = np.array(Image.open(file))\n",
        "        images.append(image)\n",
        "    \n",
        "    # Select color channel\n",
        "    color = input('What color channel are you analyzing? (red, green, blue, all)')\n",
        "    color = color.lower()\n",
        "\n",
        "    if color == 'red':\n",
        "        color = 'r'\n",
        "    elif color == 'green':\n",
        "        color = 'g'\n",
        "    elif color == 'blue':\n",
        "        color = 'b'\n",
        "    elif color == 'all':\n",
        "        pass\n",
        "    else:\n",
        "        raise Exception(f\"color should be red(r), green(g), blue(b), or all. The value of color was {chann_color}\")\n",
        "    print(color)\n",
        "\n",
        "\n",
        "    # set threshold\n",
        "    thresh = int(input(\"Select starting threshold for creating white and black image (cells will be white, range 0-255)\"))\n",
        "\n",
        "    if len(images[0].shape) == 3:\n",
        "        \n",
        "        #If red channel, select red images\n",
        "        if color == 'r':\n",
        "            image = images[0][:,:,2]\n",
        "            (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "        #If green channel, select green images\n",
        "        elif color == 'g':\n",
        "            image = image[0][:,:,1]\n",
        "            (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "        #If blue channel, select blue images\n",
        "        elif color == 'b':\n",
        "            image = images[0][:,:,0]\n",
        "            (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "        #If all channels, select all images\n",
        "        elif color == 'all':\n",
        "            image = images[0][:,:]\n",
        "            (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "        else:\n",
        "            raise Exception(f\"color should be red(r), green(g), blue(b), or all. The value of color was {color}\")\n",
        "            \n",
        "    elif len(images[0].shape) < 3:\n",
        "        image = images[0][:,:]\n",
        "        (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "        \n",
        "    else: \n",
        "        raise Exception(f\"Something is wrong with the shape of the data\")\n",
        "        \n",
        "    fig = plt.figure(figsize=(8,8))\n",
        "    plt.imshow(bw, plt.get_cmap('gray'))\n",
        "    plt.show()\n",
        "    plt.close(fig)\n",
        "    ok = input('Is this threshold OK? (yes or no)')\n",
        "    ok = ok.lower()\n",
        "\n",
        "    \n",
        "    while ok != 'yes':\n",
        "        try:\n",
        "            # set threshold\n",
        "            thresh = int(input(\"Select starting threshold for creating white and black image (cells will be white, range 0-255)\"))\n",
        "\n",
        "            if len(images[0].shape) == 3:\n",
        "                \n",
        "                #If red channel, select red images\n",
        "                if color == 'r':\n",
        "                    image = images[0][:,:,2]\n",
        "                    (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "                #If red channel, select red images\n",
        "                elif color == 'g':\n",
        "                    image = image[0][:,:,1]\n",
        "                    (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "                #If blue channel, select red images\n",
        "                elif color == 'b':\n",
        "                    image = images[0][:,:,0]\n",
        "                    (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "                #If all channels, select red images\n",
        "                elif color == 'all':\n",
        "                    image = images[0][:,:]\n",
        "                    (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                    \n",
        "            if len(images[0].shape) < 3:\n",
        "                image = images[0][:,:]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "\n",
        "            fig = plt.figure(figsize=(8,8))\n",
        "            plt.imshow(bw, plt.get_cmap('gray'))\n",
        "            plt.show()\n",
        "            plt.close(fig)\n",
        "\n",
        "            ok = input('Is this threshold OK? (yes or no)')\n",
        "            ok = ok.lower()\n",
        "\n",
        "        except:\n",
        "            raise Exception(\"Something went wrong while re-trying threshold\")\n",
        "    \n",
        "    #If yes, go through all images and append\n",
        "    if len(images[0].shape) == 3:\n",
        "        for i in range(0,len(images)):\n",
        "            #If red channel, select red images\n",
        "            if color == 'r':\n",
        "                image = images[i][:,:,2]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "\n",
        "            #If red channel, select red images\n",
        "            elif color == 'g':\n",
        "                image = images[i][:,:,1]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "\n",
        "            #If blue channel, select red images\n",
        "            elif color == 'b':\n",
        "                image = images[i][:,:,0]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "\n",
        "            #If all channels, select red images\n",
        "            elif color == 'all':\n",
        "                image = images[i][:,:]\n",
        "                (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "                image_list.append(bw)\n",
        "                    \n",
        "            else:\n",
        "                raise Exception(\"Error in batch processing color channels\")\n",
        "                    \n",
        "                    \n",
        "    elif len(image[0].shape) < 3:\n",
        "        for i in range(0,len(images)):\n",
        "            image = images[i][:,:]\n",
        "            (_, bw) = cv2.threshold(image, thresh, 255, cv2.THRESH_BINARY)\n",
        "            image_list.append(bw)\n",
        "                \n",
        "    else:\n",
        "        raise Exception('Error in confirming shape for batch processing')\n",
        "    \n",
        "    image_series = input('Would you like to see the first 9 frames? :').lower()\n",
        "    \n",
        "    if image_series == 'yes':\n",
        "        # Plot subplots of figures\n",
        "        plt.figure(figsize=(10,10))\n",
        "        for i in range(9):\n",
        "            # define subplot\n",
        "            plt.subplot(330 + 1 + i)\n",
        "            # convert to unsigned integers for viewing\n",
        "            image = image_list[i].astype('uint8')\n",
        "            # plot raw pixel data\n",
        "            plt.imshow(image, plt.get_cmap('gray'))\n",
        "            \n",
        "\n",
        "\n",
        "#Annotate images\n",
        "mass = 0\n",
        "diam = 21\n",
        "f = tp.locate(image_list[0], diameter=diam, minmass=mass, invert=True)\n",
        "\n",
        "plt.figure()  # make a new figure\n",
        "tp.annotate(f, image_list[0])\n",
        "plt.close()\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "ax.hist(f['mass'], bins=20)\n",
        "\n",
        "# Optionally, label the axes.\n",
        "ax.set(xlabel='mass', ylabel='count')\n",
        "plt.show()\n",
        "plt.close()\n",
        "\n",
        "input_ = input('Would you like to set a new minimum mass and diameter threshold? (yes or no) :').lower()\n",
        "    \n",
        "while input_ != 'no':\n",
        "    try:\n",
        "        mass, diam = map(int, input('Set minimum mass and diameter threshold (current is 21, max is 25). First number is mass, second is diameter :').split())\n",
        "        #Annotate images\n",
        "        f = tp.locate(image_list[0], diameter=diam, minmass=mass, invert=True)\n",
        "        #Plot annotation\n",
        "        plt.figure()  # make a new figure\n",
        "        tp.annotate(f, image_list[0]);\n",
        "        plt.close()\n",
        "        input_ = input('Would you like to set a new minimum mass? (yes or no) :').lower()\n",
        "            \n",
        "    except:\n",
        "        print(\"Something went wrong while re-trying threshold. Make sure to type in both numbers\")\n",
        "        continue   \n",
        "\n",
        "#Annotate all images\n",
        "\n",
        "f = tp.batch(image_list, diameter=diam, minmass=mass)\n",
        "\n",
        "#Link the trajectories \n",
        "\n",
        "linking = 'fail'\n",
        "while linking != 'success':\n",
        "    try:\n",
        "        search_range = int(input('Please indicate the maximum distance a cell can travel between frames'))\n",
        "        t = tp.link_df(f, search_range, memory=5)\n",
        "        linking = 'success'\n",
        "    except:\n",
        "        print('Something went wrong. Search range may be too big')\n",
        "\n",
        "print('Filter trajectories with spurious movement')\n",
        "min_traj = int(input('Minimum number of frames to quantify trajectory for a cell :'))\n",
        "t1 = tp.filter_stubs(t, min_traj)\n",
        "# Compare the number of particles in the unfiltered and filtered data.\n",
        "print('Before:', t['particle'].nunique())\n",
        "print('After:', t1['particle'].nunique())\n",
        "\n",
        "print('Plot trajectories')\n",
        "plt.figure(figsize=(6,6))\n",
        "t1 = t1.reset_index(drop='True')\n",
        "tp.plot_traj(t1)\n",
        "plt.close()\n",
        "\n",
        "plt.figure(figsize=(6,6))\n",
        "d = tp.compute_drift(t1)\n",
        "d.plot()\n",
        "plt.close()\n",
        "\n",
        "plt.figure(figsize=(6,6))\n",
        "tm = tp.subtract_drift(t1.copy(), d)\n",
        "ax = tp.plot_traj(tm)\n",
        "plt.show()\n",
        "plt.close\n",
        "\n",
        "\n",
        "#Initiate speed list\n",
        "speed = {}\n",
        "frame = {}\n",
        "\n",
        "#Calculate seconds between frames\n",
        "seconds = (120/len(image_list))*60\n",
        "\n",
        "for particle in np.unique(tm['particle']):\n",
        "    subset = tm[tm['particle'] == particle].reset_index(drop=True)\n",
        "    subset = subset.sort_values('frame')\n",
        "    subset = subset.set_index('frame').T\n",
        "   \n",
        "    speed_list = []\n",
        "    frame_list = []\n",
        "    for i in range(1,len(image_list)):\n",
        "        try:\n",
        "            dist=np.sqrt(np.sum((subset[:2][i] - subset[:2][i-1])**2))\n",
        "            s = (dist/seconds)*60\n",
        "            speed_list.append(s)\n",
        "            frame_list.append(i)\n",
        "        except Exception:\n",
        "            pass\n",
        "           \n",
        "    speed[f'{particle}'] = speed_list\n",
        "    frame[f'{particle}'] = frame_list\n",
        "            \n",
        "speeds = []\n",
        "frames = []\n",
        "particles = []\n",
        "\n",
        "for i in speed.keys():\n",
        "    speeds += list(speed[i])\n",
        "    frames += list(frame[i])\n",
        "    for j in range(0,len(speed[i])):\n",
        "        particles.append(i)\n",
        "        \n",
        "# Dataframe of scaled coordinates \n",
        "speed_df = pd.DataFrame.from_dict({'u_min': speeds, 'frame':frames, 'particle':particles}) \n",
        "\n",
        "#Ask for starting threshold\n",
        "\"\"\"\n",
        "# Filter total distance\n",
        "dist_threshold = int(input('Input distance traveled threshold'))\n",
        "abv_thres = []\n",
        "distances = []\n",
        "for i in tm['particle'].unique():\n",
        "    subset = tm.loc[tm['particle']== i]\n",
        "    min_frame = np.min(subset['frame'])\n",
        "    max_frame = np.max(subset['frame'])\n",
        "    x_min = float(subset.loc[subset['frame']==min_frame,'x'])\n",
        "    x_max = float(subset.loc[subset['frame']==max_frame,'x'])\n",
        "    y_min = float(subset.loc[subset['frame']==min_frame,'y'])\n",
        "    y_max = float(subset.loc[subset['frame']==max_frame,'y'])\n",
        "    dx=np.square(x_max-x_min)\n",
        "    dy=np.square(y_max-y_min)\n",
        "    dist=np.sqrt(dx+dy)\n",
        "    \n",
        "    if dist >= dist_threshold:\n",
        "        abv_thres.append(i)\n",
        "        distances.append(dist)\n",
        "        \n",
        "tm_filtered = tm[tm.particle.isin(abv_thres)]\n",
        "tm_filtered['particle'].nunique()\n",
        "\"\"\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wjIhRKFgD5Ej"
      },
      "source": [
        "# Filter total distance\n",
        "dist_threshold = int(input('Input distance traveled threshold'))\n",
        "abv_thres = []\n",
        "distances = []\n",
        "for i in tm['particle'].unique():\n",
        "    subset = tm.loc[tm['particle']== i]\n",
        "    min_frame = np.min(subset['frame'])\n",
        "    max_frame = np.max(subset['frame'])\n",
        "    x_min = float(subset.loc[subset['frame']==min_frame,'x'])\n",
        "    x_max = float(subset.loc[subset['frame']==max_frame,'x'])\n",
        "    y_min = float(subset.loc[subset['frame']==min_frame,'y'])\n",
        "    y_max = float(subset.loc[subset['frame']==max_frame,'y'])\n",
        "    dx=np.square(x_max-x_min)\n",
        "    dy=np.square(y_max-y_min)\n",
        "    dist=np.sqrt(dx+dy)\n",
        "    \n",
        "    if dist >= dist_threshold:\n",
        "        abv_thres.append(i)\n",
        "        distances.append(dist)\n",
        "        \n",
        "tm_filtered = tm[tm.particle.isin(abv_thres)]\n",
        "tm_filtered['particle'].nunique()\n",
        "\n",
        "#Scale df\n",
        "\n",
        "x_new = []\n",
        "y_new = []\n",
        "frame = []\n",
        "particle_ = []\n",
        "\n",
        "for i in tm_filtered['particle'].unique():\n",
        "    subset = tm_filtered[tm_filtered['particle']== i].reset_index(drop=True)\n",
        "    subset = subset = subset.sort_values('frame').reset_index(drop=True)\n",
        "    scale_x = subset['x'][0]\n",
        "    scale_y = subset['y'][0]\n",
        "    subset['x_scaled'] = subset['x']-scale_x\n",
        "    subset['y_scaled'] = subset['y']-scale_y\n",
        "    x_new += list(subset['x_scaled'])\n",
        "    y_new += list(subset['y_scaled'])\n",
        "    frame += list(subset['frame'])\n",
        "    particle_ += list(subset['particle'])\n",
        "\n",
        "# Dataframe of scaled coordinates \n",
        "tm_scaled = pd.DataFrame.from_dict({'x_scaled': x_new, 'y_scaled':y_new, 'frame': frame, 'particle':particle_}) \n",
        "\n",
        "plt.figure(figsize=(10, 10))\n",
        "plt.rcParams[\"font.size\"] = 12\n",
        "for i in tm_scaled['particle'].unique():\n",
        "    subset = tm_scaled[tm_scaled['particle']== i]\n",
        "    plt.plot(subset['x_scaled'], subset['y_scaled'], '-', linewidth=.5, color='black')\n",
        "    plt.xlim(-30,30)\n",
        "    plt.ylim(-30,30)\n",
        "    \n",
        "    plt.title('Flea Beetle Trajectory', fontdict={'fontsize': 14})\n",
        "    plt.xlabel(\"Distance Traveled(mM)\")\n",
        "    plt.ylabel(\"Distance Traveled(mM)\")\n",
        "    \n",
        "    plt.grid(True)\n",
        "plt.show()\n",
        "#input_ = input('Would you like to reset the distance threshold?').lower()\n",
        "plt.close()\n",
        "while input_ != 'no':\n",
        "    try:\n",
        "        # Filter total distance\n",
        "        dist_threshold = int(input('Input distance traveled threshold'))\n",
        "        abv_thres = []\n",
        "        distances = []\n",
        "        for i in tm['particle'].unique():\n",
        "            subset = tm.loc[tm['particle']== i]\n",
        "            min_frame = np.min(subset['frame'])\n",
        "            max_frame = np.max(subset['frame'])\n",
        "            x_min = float(subset.loc[subset['frame']==min_frame,'x'])\n",
        "            x_max = float(subset.loc[subset['frame']==max_frame,'x'])\n",
        "            y_min = float(subset.loc[subset['frame']==min_frame,'y'])\n",
        "            y_max = float(subset.loc[subset['frame']==max_frame,'y'])\n",
        "            dx=np.square(x_max-x_min)\n",
        "            dy=np.square(y_max-y_min)\n",
        "            dist=np.sqrt(dx+dy)\n",
        "\n",
        "            if dist >= dist_threshold:\n",
        "                abv_thres.append(i)\n",
        "                distances.append(dist)\n",
        "\n",
        "        tm_filtered = tm[tm.particle.isin(abv_thres)]\n",
        "        tm_filtered['particle'].nunique()\n",
        "\n",
        "        #Scale df\n",
        "\n",
        "        x_new = []\n",
        "        y_new = []\n",
        "        frame = []\n",
        "        particle_ = []\n",
        "\n",
        "        for i in tm_filtered['particle'].unique():\n",
        "            subset = tm_filtered[tm_filtered['particle']== i].reset_index(drop=True)\n",
        "            subset = subset = subset.sort_values('frame').reset_index(drop=True)\n",
        "            scale_x = subset['x'][0]\n",
        "            scale_y = subset['y'][0]\n",
        "            subset['x_scaled'] = subset['x']-scale_x\n",
        "            subset['y_scaled'] = subset['y']-scale_y\n",
        "            x_new += list(subset['x_scaled'])\n",
        "            y_new += list(subset['y_scaled'])\n",
        "            frame += list(subset['frame'])\n",
        "            particle_ += list(subset['particle'])\n",
        "\n",
        "        # Dataframe of scaled coordinates \n",
        "        tm_scaled = pd.DataFrame.from_dict({'x_scaled': x_new, 'y_scaled':y_new, 'frame': frame, 'particle':particle_}) \n",
        "\n",
        "        plt.figure(figsize=(10, 10))\n",
        "        plt.rcParams[\"font.size\"] = 12\n",
        "        for i in tm_scaled['particle'].unique():\n",
        "            subset = tm_scaled[tm_scaled['particle']== i]\n",
        "            plt.plot(subset['x_scaled'], subset['y_scaled'], '-', linewidth=.5, color='black')\n",
        "            plt.xlim(-30,30)\n",
        "            plt.ylim(-30,30)\n",
        "\n",
        "            plt.title('Flea Beetle Trajectory', fontdict={'fontsize': 14})\n",
        "            plt.xlabel(\"Distance Traveled(mM)\")\n",
        "            plt.ylabel(\"Distance Traveled（mM)\")\n",
        "            plt.grid(True)\n",
        "            \n",
        "        plt.show()\n",
        "        plt.close()\n",
        "        input_ = input('Would you like to reset the distance threshold?').lower()\n",
        "        break\n",
        "        \n",
        "    except:\n",
        "        print('Something went wrong. Try again')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EClBR6adD5Em"
      },
      "source": [
        "# Convert column to int for filtering\n",
        "speed_df['particle'] = speed_df['particle'].astype(int)\n",
        "\n",
        "#Filter speed df by particles\n",
        "speed_df_filt = speed_df[speed_df.particle.isin(abv_thres)]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X3p7lNHJD5Eo"
      },
      "source": [
        "# Import movie and get frames\n",
        "mov_list = []\n",
        "vidcap = cv2.VideoCapture('/content/drive/MyDrive/track/ck/CK.mp4') # Define specific folder\n",
        "success,image = vidcap.read()\n",
        "count = 0\n",
        "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "mov_list.append(image)\n",
        "while success:\n",
        "    try:\n",
        "        success,image = vidcap.read()\n",
        "        count += 1\n",
        "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "        mov_list.append(image)\n",
        "    except:\n",
        "        print(f'There is no array for image {count}')\n",
        "    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "naLenORvD5En"
      },
      "source": [
        "#Make gif of images\n",
        "plt.rcParams[\"font.size\"] = 12\n",
        "plt.rcParams[\"font.family\"] = \"arial\"\n",
        "\n",
        "for frame in range(0,tm_scaled['frame'].nunique()):\n",
        "    \n",
        "    fig = plt.figure(figsize=(9,7))\n",
        "    gs = fig.add_gridspec(nrows=2, ncols=2,wspace=0.3, hspace=0.3, width_ratios=[2, 2])\n",
        "    ax1 = fig.add_subplot(gs[0,0])\n",
        "    ax2 = fig.add_subplot(gs[0,1])\n",
        "    ax3 = fig.add_subplot(gs[1,:])\n",
        "    fig.suptitle('Flea beetle Trajectory Measurement', \n",
        "                 fontsize=14, \n",
        "                 fontweight='bold', \n",
        "                 va='top', \n",
        "                 ha='center',\n",
        "                 y = 0.96,\n",
        "                 x = 0.45)\n",
        "    \n",
        "    t2 = tm_scaled[tm_scaled['frame'] <= frame]\n",
        "    speed_subset = speed_df_filt[speed_df_filt['frame'] <= frame]\n",
        "    avg_speed = np.round(np.average(speed_subset['u_min']),2)\n",
        "    \n",
        "    ax1.imshow(mov_list[frame])\n",
        "    ax1.axis('off')\n",
        "    ax1.set_title('Raw Video', fontdict={'fontsize': 12, 'fontweight': 'bold'}, loc='left')\n",
        "    \n",
        "    for particle in t2['particle'].unique():\n",
        "        subset = t2[t2['particle'] == particle]\n",
        "            \n",
        "        if subset['particle'].shape[0] != 0:\n",
        "            ax2.plot(subset['x_scaled'], subset['y_scaled'], '-', linewidth=.5, color='black')\n",
        "            ax2.set_xlim(-250,250)\n",
        "            ax2.set_ylim(-250,250)\n",
        "    \n",
        "            ax2.set_title('Flea beetle Trajectory', fontdict={'fontsize': 12, 'fontweight': 'bold'}, loc='left')\n",
        "            ax2.set_xlabel(\"Distance Traveled X -mM\")\n",
        "            ax2.set_ylabel(\"Distance Traveled Y -mM\")\n",
        "    \n",
        "            ax2.grid(True)\n",
        "            #plt.tight_layout()\n",
        "        else:\n",
        "            pass\n",
        "        \n",
        "    sns.lineplot(x=speed_subset['frame'], y=speed_subset['u_min'], ax=ax3, color='#24a0f9')\n",
        "    ax3.set_ylabel('Average Speed--mM per min)')\n",
        "    ax3.set_xlim(0,600)\n",
        "    ax3.set_ylim(0,15)\n",
        "    ax3.set_xlabel('Frame')\n",
        "    ax3.set_title('Average Speed', fontdict={'fontsize': 12, 'fontweight': 'bold'}, loc='left')\n",
        "    \n",
        "    if frame % 50 ==0:\n",
        "        print(f\"figure {frame} completed\") \n",
        "    plt.savefig(f\"/content/drive/MyDrive/track/ck/gif/image_{frame}.png\", dpi=100)\n",
        "    plt.close()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Fq_3bhQD5Ep"
      },
      "source": [
        "#Create movie\n",
        "import re\n",
        "import cv2\n",
        "import numpy as np\n",
        "import glob\n",
        "\n",
        "\n",
        "img_array = []\n",
        "file_name = '/content/drive/MyDrive/track/ck/gif/image_'\n",
        "file_num = []\n",
        "files = []\n",
        "for filename in glob.glob('/content/drive/MyDrive/track/ck/gif/*.png'):\n",
        "    #extract image number\n",
        "    num = int(re.search(r'\\d+', filename).group())\n",
        "    file_num.append(num)\n",
        "    file_num.sort()\n",
        "\n",
        "for num in file_num:\n",
        "    file = file_name + str(num) + '.png'\n",
        "    files.append(file)\n",
        "\n",
        "for filename in files:\n",
        "    img = cv2.imread(filename)\n",
        "    height, width, layers = img.shape\n",
        "    size = (width,height)\n",
        "    img_array.append(img)\n",
        "    \n",
        "out = cv2.VideoWriter('/content/drive/MyDrive/track/ck/CKout.mp4',cv2.VideoWriter_fourcc(*'DIVX'), 19, (640, 480))\n",
        " \n",
        "for i in range(len(img_array)):\n",
        "    out.write(img_array[i])\n",
        "    \n",
        "out.release()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mhEvFb_4D5Eq"
      },
      "source": [
        "#Create gif\n",
        "import os\n",
        "import imageio as io\n",
        "images = []\n",
        "path = '/content/drive/MyDrive/track/ck/gif/'\n",
        "gif_name = 'ckgif'\n",
        "for i in range (0,len(glob.glob1(path, '*.png'))):\n",
        "    if i < 10:\n",
        "        string = f'image_{i}.png'\n",
        "        file = os.path.join(path, string)\n",
        "        image = io.imread(file)\n",
        "        images.append(image)\n",
        "    elif i < 100:\n",
        "        string = f'image_{i}.png'\n",
        "        file = os.path.join(path, string)\n",
        "        image = io.imread(file)\n",
        "        images.append(image)\n",
        "    else:\n",
        "        string = f'image_{i}.png'\n",
        "        file = os.path.join(path, string)\n",
        "        image = io.imread(file)\n",
        "        images.append(image)\n",
        "\n",
        "io.mimsave(f'/content/drive/MyDrive/track/ck/{gif_name}.gif',images,fps=19)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}